use crate::rtmp::chunk::define::CHUNK_SIZE;
use crate::rtmp::chunk::errors::PackError;
use crate::rtmp::chunk::{
    ChunkBasicHeader, ChunkHeader, ChunkInfo, ChunkMessageHeader, ExtendTimestampType,
};
use byteorder::{BigEndian, LittleEndian};
use bytesio::bytes_writer::AsyncBytesWriter;
use bytesio::bytesio::TNetIO;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;

// 解包结果枚举类型
#[derive(PartialEq, Eq, Debug)]
pub enum PackResult {
    Success,
    NotEnoughData,
}

// 块打包结构体
pub struct ChunkPacketizer {
    // 块流 ID 到块头的映射
    csid_2_chunk_header: HashMap<u32, ChunkHeader>,
    // 最大块大小
    max_chunk_size: usize,
    // 异步字节写入器
    writer: AsyncBytesWriter,
    // 扩展时间戳
    extended_timestamp: Option<u32>,
}

impl ChunkPacketizer {
    pub fn new(io: Arc<Mutex<Box<dyn TNetIO + Send + Sync>>>) -> Self {
        Self {
            csid_2_chunk_header: HashMap::new(),
            writer: AsyncBytesWriter::new(io),
            max_chunk_size: CHUNK_SIZE as usize,
            extended_timestamp: None,
        }
    }

    // 压缩块头部信息
    fn zip_chunk_header(&mut self, chunk_info: &mut ChunkInfo) -> Result<PackResult, PackError> {
        chunk_info.basic_header.format = 0;

        // 如果之前存在一个块头部信息
        if let Some(pre_header) = self
            .csid_2_chunk_header
            .get_mut(&chunk_info.basic_header.chunk_stream_id)
        {
            let cur_msg_header = &mut chunk_info.message_header;
            let pre_msg_header = &mut pre_header.message_header;

            // 如果当前时间戳小于前一个块的时间戳，发出警告
            if cur_msg_header.timestamp < pre_msg_header.timestamp {
                log::warn!(
                    "Chunk stream id: {}, the current timestamp:{}  is smaller than pre chunk timestamp: {}",
                    chunk_info.basic_header.chunk_stream_id,
                    cur_msg_header.timestamp,
                    pre_msg_header.timestamp
                );
            } else if cur_msg_header.msg_streamd_id == pre_msg_header.msg_streamd_id {
                // 如果消息流ID相同，则设置块格式为1
                chunk_info.basic_header.format = 1;
                cur_msg_header.timestamp_delta =
                    cur_msg_header.timestamp - pre_msg_header.timestamp;

                // 如果消息类型ID和消息长度相同，则设置块格式为2
                if cur_msg_header.msg_type_id == pre_msg_header.msg_type_id
                    && cur_msg_header.msg_length == pre_msg_header.msg_length
                {
                    chunk_info.basic_header.format = 2;
                    // 如果时间戳增量相同，则设置块格式为3
                    if cur_msg_header.timestamp_delta == pre_msg_header.timestamp_delta {
                        chunk_info.basic_header.format = 3;
                    }
                }
            }
        } else {
            assert_eq!(chunk_info.message_header.timestamp_delta, 0);
        }

        // 更新前一个头部信息
        self.csid_2_chunk_header.insert(
            chunk_info.basic_header.chunk_stream_id,
            ChunkHeader {
                basic_header: chunk_info.basic_header.clone(),
                message_header: chunk_info.message_header.clone(),
            },
        );

        Ok(PackResult::Success)
    }

    // 写入块基本头
    fn write_basic_header(&mut self, fmt: u8, csid: u32) -> Result<(), PackError> {
        if csid >= 64 + 255 {
            // 写入格式和块流ID
            self.writer.write_u8(fmt << 6 | 1)?;
            // 写入块流ID
            self.writer.write_u16::<BigEndian>((csid - 64) as u16)?;
        } else if csid >= 64 {
            // 写入格式
            self.writer.write_u8(fmt << 6)?;
            // 写入块流ID
            self.writer.write_u8((csid - 64) as u8)?;
        } else {
            // 写入格式和块流ID
            self.writer.write_u8(fmt << 6 | csid as u8)?;
        }

        Ok(())
    }

    // 写入块消息头部
    fn write_message_header(
        &mut self,
        basic_header: &ChunkBasicHeader,
        message_header: &mut ChunkMessageHeader,
    ) -> Result<(), PackError> {
        let message_header_timestamp: u32;
        (self.extended_timestamp, message_header_timestamp) = match basic_header.format {
            0 => {
                // 如果时间戳大于等于0xFFFFFF，则设置扩展时间戳类型为FORMAT0
                if message_header.timestamp >= 0xFFFFFF {
                    message_header.extended_timestamp_type = ExtendTimestampType::FORMAT0;
                    (Some(message_header.timestamp), 0xFFFFFF)
                } else {
                    (None, message_header.timestamp)
                }
            }
            1 | 2 => {
                // 如果时间戳增量大于等于0xFFFFFF，则设置扩展时间戳类型为FORMAT12
                if message_header.timestamp_delta >= 0xFFFFFF {
                    log::warn!(
                        "Now use extended timestamp for format {}, the value is: {}",
                        basic_header.format,
                        message_header.timestamp_delta
                    );
                    message_header.extended_timestamp_type = ExtendTimestampType::FORMAT12;
                    (Some(message_header.timestamp_delta), 0xFFFFFF)
                } else {
                    (None, message_header.timestamp_delta)
                }
            }
            _ => {
                // 不应该到达这里
                (None, 0)
            }
        };

        // 根据块格式写入消息头部
        match basic_header.format {
            0 => {
                // 写入时间戳
                self.writer
                    .write_u24::<BigEndian>(message_header_timestamp)?;
                // 写入消息长度
                self.writer
                    .write_u24::<BigEndian>(message_header.msg_length)?;
                // 写入消息类型ID
                self.writer.write_u8(message_header.msg_type_id)?;
                // 写入消息流ID
                self.writer
                    .write_u32::<LittleEndian>(message_header.msg_streamd_id)?;
            }
            1 => {
                // 写入时间戳增量
                self.writer
                    .write_u24::<BigEndian>(message_header_timestamp)?;
                // 写入消息长度
                self.writer
                    .write_u24::<BigEndian>(message_header.msg_length)?;
                // 写入消息类型ID
                self.writer.write_u8(message_header.msg_type_id)?;
            }
            2 => {
                // 写入时间戳增量
                self.writer
                    .write_u24::<BigEndian>(message_header_timestamp)?;
            }
            _ => {}
        }

        Ok(())
    }

    // 写入扩展时间戳
    fn write_extened_timestamp(&mut self, timestamp: u32) -> Result<(), PackError> {
        self.writer.write_u32::<BigEndian>(timestamp)?;

        Ok(())
    }

    // 写入块
    pub async fn write_chunk(&mut self, chunk_info: &mut ChunkInfo) -> Result<(), PackError> {
        // 压缩块头部信息
        self.zip_chunk_header(chunk_info)?;

        log::trace!(
            "write_chunk  current timestamp: {}",
            chunk_info.message_header.timestamp,
        );

        // 获取整个负载的大小
        let mut whole_payload_size = chunk_info.payload.len();
        // 写入块基本头部
        self.write_basic_header(
            chunk_info.basic_header.format,
            chunk_info.basic_header.chunk_stream_id,
        )?;
        // 写入块消息头部
        self.write_message_header(&chunk_info.basic_header, &mut chunk_info.message_header)?;
        // 写入扩展时间戳
        if let Some(extended_timestamp) = self.extended_timestamp {
            self.write_extened_timestamp(extended_timestamp)?;
        }

        let mut cur_payload_size: usize;
        while whole_payload_size > 0 {
            cur_payload_size = if whole_payload_size > self.max_chunk_size {
                self.max_chunk_size
            } else {
                whole_payload_size
            };
            // 分割负载数据
            let payload_bytes = chunk_info.payload.split_to(cur_payload_size);
            // 写入负载数据
            self.writer.write(&payload_bytes[0..])?;

            whole_payload_size -= cur_payload_size;

            if whole_payload_size > 0 {
                // 写入块基本头部
                self.write_basic_header(3, chunk_info.basic_header.chunk_stream_id)?;
                // 写入扩展时间戳
                if let Some(extended_timestamp) = self.extended_timestamp {
                    self.write_extened_timestamp(extended_timestamp)?;
                }
            }
        }
        self.writer.flush().await?;

        Ok(())
    }
}
